import torch
import torchvision
from torch import nn
import func as f
from setting import input_size, batch_size, save_path, lr, n_epoch
import ssl

ssl._create_default_https_context = ssl._create_unverified_context

# 1、数据导入
# 获取文件数据并转换成参数集
transform_train, train_set, train_loader, transform_val, val_set, val_loader = f.get_data(input_size, batch_size)

print('映射关系:', train_set.class_to_idx)  # {'cat': 0, 'dog': 1}
print('训练集长度:', len(train_set.imgs))  # 22500
print('训练集规格:', train_set[1][0].size())  # torch.Size([3, 224, 224])

device = torch.device('cpu')  # 选择训练模式(CPU)
print('训练模式:', device, '模式')

# 残差网络(18指定的是带有权重的18层，包括卷积层和全连接层，不包括池化层和BN层)
# 2、导入resnet18模型
# 数据导入后，就可以作为模型的输入了。这里使用的是PyTorch分类中的ResNet网络结构
# 预训练网络一般是在1000类的ImageNet数据集上进行的，迁移到自己数据集的二分类，需要替换最后的全连接层为所需要的输出
# 下面这三行代码进行的就是用models模块导入resnet18网络，然后获取全连接层的输入channel个数，用这个channel个数和要做的分类类别数（这里是2）替换原来模型中的全连接层

# 用torchvision.models模块来导入模型
transfer_model = torchvision.models.resnet18(pretrained=True)  # pretrained=True 使用预训练模型
# 获取全连接层的输入 channel 个数
num_ftrs = transfer_model.fc.in_features
# 用这个channel个数和要做的分类类别数（这里是2）替换原来模型中的全连接层
transfer_model.fc = nn.Linear(num_ftrs, 2)

# 3、构建神经网络
net = transfer_model.to(device)

# 4、定义交叉熵损失函数criterion = nn.CrossEntropyLoss()
# 在PyTorch中采用torch.nn模块来定义网络的所有层，比如卷积、降采样、损失层等等，这里采用交叉熵函数
criterion = nn.CrossEntropyLoss()

# 5、定义优化函数
# SGD优化器——随机梯度下降 # 学习率lr=10^-3
# SGD优化器会把数据拆分后再分批不断放入nn中计算，每次使用批数据，虽然不能反映整体数据的情况，不过却很大程度上加速了nn的训练过程, 而且也不会丢失太多准确率.
optimizer = torch.optim.SGD(net.fc.parameters(), lr=lr)

for epoch in range(n_epoch):
    print('第{}次训练'.format(epoch + 1))
    # 6、调用func中的train.py文件进行训练
    f.train(net, optimizer, device, criterion, train_loader)

    # 7、验证训练的正确性，输出网络精度
    f.validate(net, device, val_loader)

# 仅保存和加载模型参数
torch.save(net.state_dict(), save_path)
